from app.core.config import settings
from contextlib import asynccontextmanager
from fastapi import FastAPI
from tortoise.contrib.fastapi import RegisterTortoise
from tortoise import connections
from loguru import logger

# 需要补充 model 类
TORTOISE_MODULES = []

TORTOISE_ORM = {
    "connections": {
        "mysql": {
            # 数据库连接类型
            "engine": "tortoise.backends.mysql",
            "credentials": {
                "host": settings.MYSQL_HOST,
                "port": settings.MYSQL_PORT,
                "user": settings.MYSQL_USER,
                "password": settings.MYSQL_PASSWORD,
                "database": settings.MYSQL_DB,
            },
        },
    },
    "apps": {
        "models": {
            "models": TORTOISE_MODULES,
            # 使用的配置
            "default_connection": "mysql",
        },
    },
    # 关闭默认时区
    "use_tz": False,
    # 本地时间配置
    "timezone": "Asia/Shanghai",
}


# 使用 lifespan 管理应用生命周期
@asynccontextmanager
async def init_mysql(app: FastAPI):
    try:
        async with RegisterTortoise(
            app,
            config=TORTOISE_ORM,
            generate_schemas=True,
            add_exception_handlers=True,
        ):
            yield logger.info("MySQL 数据库初始化完成")
            await connections.close_all()
            logger.info("MySQL 数据库连接已关闭")
    except Exception as e:
        logger.error(f"MySQL 数据库初始化失败: {e}")
        raise e
